import torch
from torch import nn
from torchviz import make_dot, make_dot_from_trace
import models

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
pretrained_dispnet = "/media/gzy/My Passport/Unsupervised_Monocular_Depth_Estimation/checkpoint/resnet_laplace/09-27-15:26/dispnet_model_best.pth.tar"

dispnet = "DispResNet"
disp_net = getattr(models, dispnet)(16).to(device) if dispnet in ['DispSENet', 'DispDenseNet'] else getattr(models, dispnet)().to( device)
weights = torch.load(pretrained_dispnet, map_location=torch.device('cpu'))
disp_net.load_state_dict(weights['state_dict'])
disp_net.eval()

x = torch.randn(1, 3, 256, 832)
y = disp_net(x)

make_dot(y.mean(), params=dict(disp_net.named_parameters()), show_attrs=True, show_saved=True).view()
